EKS オーケストレータを使った SageMaker HyperPod クラスターで FSx for Lustre をマウントしてみた

###################################################
# Lustre File System
###################################################
resource "aws_fsx_lustre_file_system" "this" {
  storage_type                = "SSD"
  file_system_type_version    = "2.15"
  storage_capacity            = 1200
  security_group_ids          = [aws_security_group.lustre.id]
  subnet_ids                  = [module.vpc.private_subnets[0]]
  data_compression_type       = "LZ4"
  deployment_type             = "PERSISTENT_2"
  per_unit_storage_throughput = 250

  metadata_configuration {
    mode = "AUTOMATIC"
  }
}

resource "aws_fsx_data_repository_association" "this" {
  file_system_id       = aws_fsx_lustre_file_system.this.id
  data_repository_path = "s3://${aws_s3_bucket.data_repository.bucket}"
  file_system_path     = "/"

  s3 {
    auto_export_policy {
      events = ["NEW", "CHANGED", "DELETED"]
    }

    auto_import_policy {
      events = ["NEW", "CHANGED", "DELETED"]
    }
  }
}

###################################################
# Security Group for Lustre File System
###################################################
resource "aws_security_group" "lustre" {
  name        = "${local.prefix}-lustre-sg"
  vpc_id      = module.vpc.vpc_id
  description = "${local.prefix}-hyperpod-sg"

  tags = {
    Name = "${local.prefix}-lustre-sg"
  }
}
# Ingress
resource "aws_vpc_security_group_ingress_rule" "lustre_all_traffic_self" {
  security_group_id            = aws_security_group.lustre.id
  referenced_security_group_id = aws_security_group.lustre.id
  ip_protocol                  = "-1"
}

resource "aws_vpc_security_group_ingress_rule" "lustre_all_traffic_hyperpod" {
  security_group_id            = aws_security_group.lustre.id
  referenced_security_group_id = aws_security_group.hyperpod.id
  ip_protocol                  = "-1"
}

# Egress
resource "aws_vpc_security_group_egress_rule" "lustre_all_traffic_self" {
  security_group_id            = aws_security_group.lustre.id
  referenced_security_group_id = aws_security_group.lustre.id
  ip_protocol                  = "-1"
}

resource "aws_vpc_security_group_egress_rule" "lustre_all_traffic_hyperpod" {
  security_group_id            = aws_security_group.lustre.id
  referenced_security_group_id = aws_security_group.hyperpod.id
  ip_protocol                  = "-1"
}

###################################################
# Security Group for SageMaker HyperPod Cluster
###################################################
resource "aws_security_group" "hyperpod" {
  name        = "${local.prefix}-hyperpod-sg"
  vpc_id      = module.vpc.vpc_id
  description = "${local.prefix}-hyperpod-sg"

  tags = {
    Name = "${local.prefix}-hyperpod-sg"
  }
}

resource "aws_vpc_security_group_ingress_rule" "hyperpod_allow_all_traffic_ipv4" {
  security_group_id            = aws_security_group.hyperpod.id
  referenced_security_group_id = aws_security_group.hyperpod.id
  ip_protocol                  = "-1"
}

resource "aws_vpc_security_group_ingress_rule" "hyperpod_allow_all_traffic_eks" {
  security_group_id            = aws_security_group.hyperpod.id
  referenced_security_group_id = aws_security_group.eks.id
  ip_protocol                  = "-1"
}

resource "aws_vpc_security_group_egress_rule" "hyperpod_allow_all_traffic_ipv4" {
  security_group_id = aws_security_group.hyperpod.id
  cidr_ipv4         = "0.0.0.0/0"
  ip_protocol       = "-1"
}

resource "aws_vpc_security_group_egress_rule" "hyperpod_allow_all_traffic_self" {
  security_group_id            = aws_security_group.hyperpod.id
  referenced_security_group_id = aws_security_group.hyperpod.id
  ip_protocol                  = "-1"
}

###################################################
# FSx CSI Driver
###################################################
data "tls_certificate" "this" {
  url = aws_eks_cluster.this.identity[0].oidc[0].issuer
  depends_on = [
    aws_eks_cluster.this
  ]
}

resource "aws_iam_openid_connect_provider" "this" {
  url = aws_eks_cluster.this.identity[0].oidc[0].issuer

  client_id_list = [
    "sts.amazonaws.com",
  ]

  thumbprint_list = [
    data.tls_certificate.this.certificates[0].sha1_fingerprint
  ]
}

data "aws_iam_policy_document" "assume_csi_driver" {
  statement {
    actions = [
      "sts:AssumeRoleWithWebIdentity",
    ]

    effect = "Allow"

    principals {
      type        = "Federated"
      identifiers = [aws_iam_openid_connect_provider.this.arn]
    }

    condition {
      test     = "StringEquals"
      variable = "${aws_iam_openid_connect_provider.this.url}:aud"
      values = [
        "sts.amazonaws.com",
      ]
    }
    condition {
      test     = "StringEquals"
      variable = "${aws_iam_openid_connect_provider.this.url}:sub"
      values = [
        "system:serviceaccount:kube-system:fsx-csi-controller-sa",
      ]
    }
  }
}

resource "aws_iam_role" "fsx_csi_driver" {
  name = "${local.prefix}-fsx-csi-driver-role"

  assume_role_policy = data.aws_iam_policy_document.assume_csi_driver.json
}

resource "aws_iam_role_policy_attachment" "fsx_csi_driver" {
  role       = aws_iam_role.fsx_csi_driver.name
  policy_arn = "arn:aws:iam::aws:policy/AmazonFSxFullAccess"
}

resource "helm_release" "aws_fsx_csi_driver" {
  name              = "aws-fsx-csi-driver"
  repository        = "https://kubernetes-sigs.github.io/aws-fsx-csi-driver"
  chart             = "aws-fsx-csi-driver"
  namespace         = "kube-system"
  dependency_update = true
  wait              = false

  depends_on = [
    aws_eks_access_policy_association.this
  ]
}

data "kubernetes_service_account_v1" "this" {
  metadata {
    name      = "fsx-csi-controller-sa"
    namespace = "kube-system"
  }

  depends_on = [
    helm_release.aws_fsx_csi_driver,
    aws_eks_access_policy_association.this
  ]
}

resource "kubernetes_annotations" "this" {
  api_version = "v1"
  kind        = "ServiceAccount"

  metadata {
    name      = data.kubernetes_service_account_v1.this.metadata[0].name
    namespace = data.kubernetes_service_account_v1.this.metadata[0].namespace
  }

  annotations = {
    "eks.amazonaws.com/role-arn" = aws_iam_role.fsx_csi_driver.arn
  }

  depends_on = [
    aws_eks_access_policy_association.this
  ]
}

###################################################
# FSx CSI Driver
###################################################
data "tls_certificate" "this" {
  url = aws_eks_cluster.this.identity[0].oidc[0].issuer
  depends_on = [
    aws_eks_cluster.this
  ]
}

resource "aws_iam_openid_connect_provider" "this" {
  url = aws_eks_cluster.this.identity[0].oidc[0].issuer

  client_id_list = [
    "sts.amazonaws.com",
  ]

  thumbprint_list = [
    data.tls_certificate.this.certificates[0].sha1_fingerprint
  ]
}

data "aws_iam_policy_document" "assume_csi_driver" {
  statement {
    actions = [
      "sts:AssumeRoleWithWebIdentity",
    ]

    effect = "Allow"

    principals {
      type        = "Federated"
      identifiers = [aws_iam_openid_connect_provider.this.arn]
    }

    condition {
      test     = "StringEquals"
      variable = "${aws_iam_openid_connect_provider.this.url}:aud"
      values = [
        "sts.amazonaws.com",
      ]
    }
    condition {
      test     = "StringEquals"
      variable = "${aws_iam_openid_connect_provider.this.url}:sub"
      values = [
        "system:serviceaccount:kube-system:fsx-csi-controller-sa",
      ]
    }
  }
}

resource "aws_iam_role" "fsx_csi_driver" {
  name = "${local.prefix}-fsx-csi-driver-role"

  assume_role_policy = data.aws_iam_policy_document.assume_csi_driver.json
}

resource "aws_iam_role_policy_attachment" "fsx_csi_driver" {
  role       = aws_iam_role.fsx_csi_driver.name
  policy_arn = "arn:aws:iam::aws:policy/AmazonFSxFullAccess"
}

resource "helm_release" "aws_fsx_csi_driver" {
  name              = "aws-fsx-csi-driver"
  repository        = "https://kubernetes-sigs.github.io/aws-fsx-csi-driver"
  chart             = "aws-fsx-csi-driver"
  namespace         = "kube-system"
  dependency_update = true
  wait              = false

  depends_on = [
    aws_eks_access_policy_association.this
  ]
}

data "kubernetes_service_account_v1" "this" {
  metadata {
    name      = "fsx-csi-controller-sa"
    namespace = "kube-system"
  }

  depends_on = [
    helm_release.aws_fsx_csi_driver,
    aws_eks_access_policy_association.this
  ]
}

resource "kubernetes_annotations" "this" {
  api_version = "v1"
  kind        = "ServiceAccount"

  metadata {
    name      = data.kubernetes_service_account_v1.this.metadata[0].name
    namespace = data.kubernetes_service_account_v1.this.metadata[0].namespace
  }

  annotations = {
    "eks.amazonaws.com/role-arn" = aws_iam_role.fsx_csi_driver.arn
  }

  depends_on = [
    aws_eks_access_policy_association.this
  ]
}

###################################################
# Persistent Volume and Persistent Volume Claim
###################################################
resource "kubernetes_storage_class_v1" "this" {
  metadata {
    name = "fsx-sc"
  }
  storage_provisioner = "fsx.csi.aws.com"
  parameters = {
    "fileSystemId"     = aws_fsx_lustre_file_system.this.id
    "subnetId"         = aws_fsx_lustre_file_system.this.subnet_ids[0]
    "securityGroupIds" = aws_security_group.lustre.id
  }

  depends_on = [
    aws_eks_access_policy_association.this
  ]
}

resource "kubernetes_persistent_volume_v1" "this" {
  metadata {
    name = "fsx-pv"
  }

  spec {
    capacity = {
      storage = "1200Gi"
    }
    access_modes                     = ["ReadWriteMany"]
    volume_mode                      = "Filesystem"
    persistent_volume_reclaim_policy = "Retain"
    storage_class_name               = kubernetes_storage_class_v1.this.metadata[0].name
    persistent_volume_source {
      csi {
        driver        = "fsx.csi.aws.com"
        volume_handle = aws_fsx_lustre_file_system.this.id
        volume_attributes = {
          dnsname   = aws_fsx_lustre_file_system.this.dns_name
          mountname = aws_fsx_lustre_file_system.this.mount_name
        }
      }
    }
  }

  depends_on = [
    aws_eks_access_policy_association.this
  ]
}

resource "kubernetes_persistent_volume_claim_v1" "this" {
  metadata {
    name = "fsx-claim"
  }
  spec {
    access_modes       = ["ReadWriteMany"]
    storage_class_name = kubernetes_storage_class_v1.this.metadata[0].name
    resources {
      requests = {
        storage = "1200Gi"
      }
    }
  }

  depends_on = [
    aws_eks_access_policy_association.this,
    kubernetes_persistent_volume_v1.this
  ]
}

cat <<EOF> pv.yaml
apiVersion: v1
kind: PersistentVolume
metadata:
  name: fsx-pv
spec:
  capacity:
    storage: 1200Gi  # Adjust based on your FSx volume size
  volumeMode: Filesystem
  accessModes:
    - ReadWriteMany
  persistentVolumeReclaimPolicy: Retain
  storageClassName: fsx-sc
  csi:
    driver: fsx.csi.aws.com
    volumeHandle: fs-xxxxx  # Replace with your FSx file system ID
EOF

kubectl apply -f pv.yaml

MountVolume.SetUp failed for volume "fsx-pv" : rpc error: code = InvalidArgument desc = dnsname is not provided

###################################################
# Hello Sample file
###################################################
resource "aws_s3_object" "hello_from_s3" {
  bucket = aws_s3_bucket.data_repository.bucket
  key    = "hello_from_s3.txt"

  content = "hello! from S3!"

  depends_on = [aws_fsx_data_repository_association.this]
}

############################################
# Sample Kubernetes Pod
############################################
resource "kubernetes_pod_v1" "this" {
  metadata {
    name = "fsx-app"
  }
  spec {
    container {
      name    = "app"
      image   = "centos"
      command = ["/bin/sh"]
      args    = ["-c", "while true; do echo $(date -u) >> /data/out.txt; sleep 5; done"]
      volume_mount {
        name       = "persistent-storage"
        mount_path = "/data"
      }
    }
    volume {
      name = "persistent-storage"
      persistent_volume_claim {
        claim_name = kubernetes_persistent_volume_claim_v1.this.metadata[0].name
      }
    }
  }

  depends_on = [
    kubernetes_persistent_volume_v1.this,
    awscc_sagemaker_cluster.this
  ]
}

takakuni@ sagemaker_hyperpod_lustre_eks % kubectl exec -it fsx-app -- /bin/bash
[root@fsx-app /]# df -h
Filesystem                Size  Used Avail Use% Mounted on
overlay                   500G  6.8G  494G   2% /
tmpfs                      64M     0   64M   0% /dev
tmpfs                     7.8G     0  7.8G   0% /sys/fs/cgroup
10.0.1.159@tcp:/jm2njbev  1.2T  7.5M  1.2T   1% /data
/dev/nvme0n1p1            100G   20G   81G  20% /etc/hosts
/dev/nvme1n1              500G  6.8G  494G   2% /etc/hostname
shm                        64M     0   64M   0% /dev/shm
tmpfs                      15G   12K   15G   1% /run/secrets/kubernetes.io/serviceaccount
tmpfs                     7.8G     0  7.8G   0% /proc/acpi
tmpfs                     7.8G     0  7.8G   0% /sys/firmware

[root@fsx-app /]# cat /data/
hello_from_s3.txt  out.txt
[root@fsx-app /]# cat /data/hello_from_s3.txt
hello! from S3![root@fsx-app /]#